import numpy as np
from keras.applications.vgg16 import VGG16
from keras.applications.vgg16 import preprocess_input as preprocess_input_vgg
from keras.preprocessing import image
from numpy import linalg as LA
import os,h5py
class VGGNet():
    def __init__(self):
        self.input_shape=(224,224,3)
        self.weight='imagenet'
        self.pooling='max'
        self.model_vgg=VGG16(weights=self.weight,
                             input_shape=(self.input_shape[0],self.input_shape[1],self.input_shape[2]),
                             pooling=self.pooling,include_top=False
                             )
        self.model_vgg.predict(np.zeros((1,224,224,3)))

    def vgg_extract_feat(self,img_path):
        img=image.load_img(img_path,target_size=(self.input_shape[0],self.input_shape[1]))
        img=image.img_to_array(img)
        img=np.expand_dims(img,axis=0)
        img=preprocess_input_vgg(img)
        feat=self.model_vgg.predict(img)
        norm_feat=feat[0] / LA.norm(feat[0])
        return  norm_feat
def get_imlist(path):
    return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.png') or f.endswith('.jpg')]

if __name__=="__main__":
    database="img"
    index='1.h5'
    img_list=get_imlist(database)
    print ("--------------------------")
    feats=[]
    names=[]
    model=VGGNet()
    for i,image_path in enumerate(img_list):
        norm_feat=model.vgg_extract_feat(image_path)
        img_name=os.path.split(image_path)[1]
        feats.append(norm_feat)
        names.append(img_name.encode())
        print ("extractingg feature from image No")
    feats=np.array(feats)
    output=index
    print ("--------------------------111")
    h5f=h5py.File(output,'w')
    h5f.create_dataset('dataset_1',data=feats)
    h5f.create_dataset('dataset_2',data=np.string_(names))
    h5f.close()
